"""
数据存储模块
将爬取的文献存储到SQLite数据库
"""
import sqlite3
import json
from pathlib import Path
from typing import List, Dict, Optional
import logging

from config import DATA_DIR, RAW_DATA_DIR

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)


class DataStorage:
    """文献数据存储"""
    
    def __init__(self, db_path: str = None):
        if db_path is None:
            db_path = DATA_DIR / "medical_literature.db"
        
        self.db_path = db_path
        self._init_database()
    
    def _init_database(self):
        """初始化数据库表结构"""
        Path(self.db_path).parent.mkdir(parents=True, exist_ok=True)
        
        conn = sqlite3.connect(self.db_path)
        cursor = conn.cursor()
        
        cursor.execute("""
            CREATE TABLE IF NOT EXISTS papers (
                pmid TEXT PRIMARY KEY,
                title TEXT NOT NULL,
                abstract TEXT,
                authors TEXT,
                journal TEXT,
                publication_date TEXT,
                keywords TEXT,
                full_text_url TEXT,
                created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
            )
        """)
        
        cursor.execute("""
            CREATE TABLE IF NOT EXISTS embeddings (
                pmid TEXT PRIMARY KEY,
                embedding_vector BLOB,
                FOREIGN KEY (pmid) REFERENCES papers(pmid)
            )
        """)
        
        conn.commit()
        conn.close()
        logger.info(f"数据库初始化完成: {self.db_path}")
    
    def save_papers(self, papers: List[Dict]) -> int:
        """
        保存论文到数据库
        
        Args:
            papers: 论文列表
            
        Returns:
            成功保存的数量
        """
        conn = sqlite3.connect(self.db_path)
        cursor = conn.cursor()
        
        saved_count = 0
        
        for paper in papers:
            try:
                cursor.execute("""
                    INSERT OR REPLACE INTO papers 
                    (pmid, title, abstract, authors, journal, publication_date, keywords, full_text_url, updated_at)
                    VALUES (?, ?, ?, ?, ?, ?, ?, ?, CURRENT_TIMESTAMP)
                """, (
                    paper.get("pmid"),
                    paper.get("title"),
                    paper.get("abstract"),
                    json.dumps(paper.get("authors", [])),
                    paper.get("journal"),
                    paper.get("publication_date"),
                    json.dumps(paper.get("keywords", [])),
                    paper.get("full_text_url")
                ))
                saved_count += 1
            except Exception as e:
                logger.error(f"保存论文失败: {e}")
        
        conn.commit()
        conn.close()
        
        logger.info(f"成功保存 {saved_count} 篇论文")
        return saved_count
    
    def get_all_papers(self, limit: int = None) -> List[Dict]:
        """获取所有论文"""
        conn = sqlite3.connect(self.db_path)
        cursor = conn.cursor()
        
        query = "SELECT * FROM papers"
        if limit:
            query += f" LIMIT {limit}"
        
        cursor.execute(query)
        rows = cursor.fetchall()
        
        papers = []
        for row in rows:
            papers.append({
                "pmid": row[0],
                "title": row[1],
                "abstract": row[2],
                "authors": json.loads(row[3]) if row[3] else [],
                "journal": row[4],
                "publication_date": row[5],
                "keywords": json.loads(row[6]) if row[6] else [],
                "full_text_url": row[7]
            })
        
        conn.close()
        return papers
    
    def get_paper_by_pmid(self, pmid: str) -> Optional[Dict]:
        """根据PMID获取论文"""
        conn = sqlite3.connect(self.db_path)
        cursor = conn.cursor()
        
        cursor.execute("SELECT * FROM papers WHERE pmid = ?", (pmid,))
        row = cursor.fetchone()
        conn.close()
        
        if row:
            return {
                "pmid": row[0],
                "title": row[1],
                "abstract": row[2],
                "authors": json.loads(row[3]) if row[3] else [],
                "journal": row[4],
                "publication_date": row[5],
                "keywords": json.loads(row[6]) if row[6] else [],
                "full_text_url": row[7]
            }
        return None
    
    def get_paper_count(self) -> int:
        """获取论文总数"""
        conn = sqlite3.connect(self.db_path)
        cursor = conn.cursor()
        cursor.execute("SELECT COUNT(*) FROM papers")
        count = cursor.fetchone()[0]
        conn.close()
        return count


if __name__ == "__main__":
    # 测试存储
    storage = DataStorage()
    print(f"数据库中已有 {storage.get_paper_count()} 篇论文")
